from db.es.es_service import ElasticsearchStorage
import pandas as pd
from datetime import datetime
import settings
import json
# 创建Elasticsearch客户端
es = ElasticsearchStorage()


def find_session_with_text(text_info):
    text_info = '我们会根据这些信息制定一个更加个性化的退休规划方案'
    query = {
        "query": {
            "match_phrase": {
                "text_info": {
                    "query": text_info
                }
            }
        }
    }
    results = es.search_documents(settings.CONVERSATION_HISTORY_INDEX, query)
    df = pd.DataFrame(columns=["对话id", "角色", "对话内容", "对话时间", "态度", "踩的原因", "用户建议"])
    if results:
        session_id = ''
        # 提取查询结果
        for hit in results:
            session_id = hit["_source"]["session_id"]

        query = {
            "query": {
                "bool": {
                    "must": [
                        {
                            "term": {
                                "session_id": {
                                    "value": session_id
                                }
                            }
                        }
                    ]
                }
            }, "sort": [
                {
                    "message_time": {
                        "order": "asc"
                    }
                }
            ]
            , "size": 1000
        }
        results = es.search_documents(settings.CONVERSATION_HISTORY_INDEX, query)
        if results:
            # 遍历查询结果
            for hit in results:
                session_id = hit["_source"]["session_id"]
                role = hit["_source"]["role"]
                text_info = hit["_source"]["text_info"]
                message_time = hit["_source"]["message_time"]
                attitude = hit["_source"].get("attitude", None)
                badoption = hit["_source"].get("badoption", None)
                userPrefer = hit["_source"].get("userPrefer", None)
                df = df.append({"对话id": session_id, "角色": role, "对话内容": text_info, "对话时间": message_time, "态度": attitude,
                                "踩的原因": badoption,
                                "用户建议": userPrefer},
                               ignore_index=True)

        # 导出到Excel文件
        df.to_excel("output.xlsx", index=False)


def find_session_with_date():
    session_ids = set([])
    # today = datetime.date.today()

    # 设置查询时间范围：从今天早上0点到晚上12点
    # start_time = today.strftime("%Y-%m-%dT00:00:00")
    # end_time = today.strftime("%Y-%m-%dT23:59:59")

    start_time_str = "2023-07-19 01:00:00"
    start_time = datetime.strptime(start_time_str, "%Y-%m-%d %H:%M:%S")

    end_time_str = "2023-07-19 23:00:00"
    end_time = datetime.strptime(end_time_str, "%Y-%m-%d %H:%M:%S")

    # 定义查询条件xw
    query = {
  "query": {
    "range": {
      "message_time": {
        "gte": "2023-07-19T12:00:00",
        "lt": "2023-07-19T23:00:00"
      }
    }
  },
  "sort": [
    {"session_id.keyword": "asc"},
    {"message_time": "asc"}
  ],
  "size": 10000
}
    # 执行查询
    results = es.search_documents(settings.QW_CONVERSATION_HISTORY_INDEX, query)

    if results:
        # 创建一个空的DataFrame
        data = []
        # 提取查询结果
        for idx, hit in enumerate(results):
            session_id = hit["_source"]["session_id"]
            role = hit["_source"]["role"]
            detail = hit["_source"]["text_info"]
            msg_time = hit["_source"]["message_time"]
            if idx > 0 and results[idx]["_source"]["session_id"] != results[idx-1]["_source"]["session_id"]:
                data.append(["" for _ in range(5)])
            if role == 'assistant':
                data.append([session_id, role, json.loads(detail)['最终话术'], detail, msg_time])
            else:
                data.append([session_id, role, detail, '', msg_time])

        # 导出到Excel文件

        df = pd.DataFrame(columns=["session_id", "role", "content", "detail", "msg_time"], data=data)
        print(len(results), len(df), len(data))
        # df.sort_values(['session_id', 'msg_time'])
        print(len(results), len(df), len(data))
        df.to_excel("./20230719.xlsx", index=False)


if __name__ == '__main__':
    find_session_with_date()
